package com.sqs.service.impl;

import com.sqs.entity.AIRequest;
import com.sqs.entity.ChatMessage;
import com.sqs.entity.Message;
import com.sqs.enumration.AIRoleType;
import com.sqs.mapper.ChatMessageMapper;
import com.sqs.service.IChatMessageService;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import java.util.ArrayList;
import java.util.List;


@Service
public class ChatMessageServiceImpl extends ServiceImpl<ChatMessageMapper, ChatMessage> implements IChatMessageService {

    @Autowired
    private WebClient webClient;
    @Autowired
    private ChatMessageMapper chatMessageMapper;

    @Override
    public Flux<String> streamChat(Long sessionId, String question) {
        // 根据需要设置 baseUrl 和 header
        return webClient.mutate()
                .build()
                .post()
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(buildRequestBody(sessionId, question))
                .retrieve()
                .bodyToFlux(String.class)
                .filter(chunk -> !chunk.isBlank() && !chunk.equals("[DONE]"))
                .map(this::extractContentFromChunk);
    }

    /**
     * 构造ai请求体
     * */
    private AIRequest buildRequestBody(Long sessionId, String question) {
        //1.查询历史对话记录
        List<ChatMessage> chatMessages = lambdaQuery()
                .eq(ChatMessage::getSessionId, sessionId)
                .orderByAsc(ChatMessage::getCreateTime)
                .list();
        List<Message> messages = new ArrayList<>();
        //2.构造对话记录
        messages.add(new Message(AIRoleType.ASSISTANT.getRole(), "你是一个数学解答助手"));
        chatMessages.forEach(chatMessage -> {
            messages.add(new Message(AIRoleType.USER.getRole(), chatMessage.getQuestion()));
            messages.add(new Message(AIRoleType.SYSTEM.getRole(), chatMessage.getAnswer()));
        });
        messages.add(new Message(AIRoleType.USER.getRole(), question));

        //3.构造流式请求体
        return AIRequest.builder()
                .message(messages)
                .build();
    }

    private String extractContentFromChunk(String chunk) {
        // 这里只做简化：去掉 "data: " 前缀并解析 JSON，取 choices[0].delta.content
        try {
            String json = chunk.startsWith("data:") ? chunk.substring(5).trim() : chunk;
            // 解析 JSON
            com.fasterxml.jackson.databind.JsonNode node =
                    com.fasterxml.jackson.databind.json.JsonMapper.builder().build()
                            .readTree(json);
            var choices = node.path("choices");
            if (choices.isArray() && choices.size() > 0) {
                var delta = choices.get(0).path("delta");
                var contentNode = delta.path("content");
                if (!contentNode.isMissingNode()) {
                    return contentNode.asText();
                }
            }
        } catch (Exception e) {
            // 解析异常时忽略或记录
        }
        return "";
    }

}
